import torch
import torch.nn as nn
import numpy as np
from demo3 import DeepseekMLA_Prefill_Windows,PretrainedConfig


"""
从 Prefill 张量到 Token 文本输出的完整流程演示
代码说明
1. MockTokenizer 类
模拟了一个简单的分词器，将 token ID 映射到可读的 token 文本

包含了一些特殊 token（如 [PAD], [UNK] 等）

提供了 decode 方法将 token ID 列表转换为文本

2. LMHead 类
语言模型头，将隐藏状态映射到词汇表大小的 logits

这是一个简单的线性层

3. DeepseekMLA_TextGeneration 类
扩展了之前的 DeepseekMLA_Prefill_Windows 类

添加了 generate_text 方法，实现了从 prefill 输出到 token ID 的完整流程

包含了温度缩放、top-k 和 top-p 采样等文本生成技术

4. 文本生成流程
获取 prefill 输出

通过 LM Head 获取 logits

应用温度参数调整 logits

可选：应用 top-k 过滤，只保留概率最高的 k 个 token

可选：应用 top-p（核采样）过滤，只保留累积概率达到 p 的最小 token 集合

从处理后的分布中采样，得到 token ID

使用分词器将 token ID 解码为文本

5. 解码阶段演示
展示了如何从生成的最后一个 token 继续生成更多文本

在实际应用中，这会是一个循环过程，直到生成结束标记或达到最大长度

这个演示代码展示了从 prefill 张量到最终文本输出的完整流程，包括各种文本生成技术的实现。在实际应用中，您需要使用与模型训练时相同的真实分词器，而不是这个模拟的分词器。
"""

# 模拟分词器类 - 用于演示目的
class MockTokenizer:
    def __init__(self, vocab_size=10000):
        self.vocab_size = vocab_size
        # 创建一些模拟的词汇表映射
        self.id_to_token = {i: f"token_{i}" for i in range(vocab_size)}
        # 添加一些特殊token
        self.id_to_token[0] = "[PAD]"
        self.id_to_token[1] = "[UNK]"
        self.id_to_token[2] = "[CLS]"
        self.id_to_token[3] = "[SEP]"
        self.id_to_token[4] = "[MASK]"

        # 创建反向映射
        self.token_to_id = {v: k for k, v in self.id_to_token.items()}

    def decode(self, token_ids, skip_special_tokens=True):
        """将token ID列表转换为文本"""
        tokens = []
        for token_id in token_ids:
            if token_id < self.vocab_size:
                token = self.id_to_token[token_id]
                if skip_special_tokens and token.startswith("[") and token.endswith("]"):
                    continue
                tokens.append(token)
            else:
                tokens.append("[UNK]")

        return " ".join(tokens)

    def __call__(self, text):
        """简单模拟编码过程 - 实际分词器会更复杂"""
        # 这里只是简单地将文本按空格分割并映射到ID
        words = text.split()
        return {"input_ids": torch.tensor([self.token_to_id.get(word, 1) for word in words])}


# 语言模型头 (LM Head)
class LMHead(nn.Module):
    def __init__(self, hidden_size, vocab_size):
        super().__init__()
        self.linear = nn.Linear(hidden_size, vocab_size)

    def forward(self, x):
        return self.linear(x)


# 扩展之前的 DeepseekMLA_Prefill_Windows 类，添加文本生成功能
class DeepseekMLA_TextGeneration(DeepseekMLA_Prefill_Windows):
    def __init__(self, *args, vocab_size=10000, **kwargs):
        super().__init__(*args, **kwargs)
        self.vocab_size = vocab_size
        self.lm_head = LMHead(self.hidden_size, vocab_size)

    def generate_text(self, hidden_states, positions, slot_mapping, temperature=1.0, top_k=50, top_p=0.9):
        """
        从prefill输出生成文本

        参数:
            hidden_states: 输入隐藏状态 [seq_len, hidden_size]
            positions: 位置编码 [seq_len]
            slot_mapping: 缓存映射 [seq_len]
            temperature: 温度参数，控制随机性
            top_k: 只考虑概率最高的k个token
            top_p: 核采样参数，只考虑累积概率达到p的最小token集合
        """
        # 1. 获取prefill输出
        prefill_output, kv_cache = self.forward_prefill(hidden_states, positions, slot_mapping)

        # 2. 通过LM Head获取logits
        logits = self.lm_head(prefill_output)  # [seq_len, vocab_size]

        # 3. 应用温度缩放
        logits = logits / temperature

        # 4. Top-k过滤
        if top_k > 0:
            top_k_values, top_k_indices = torch.topk(logits, top_k, dim=-1)
            min_values = top_k_values[:, -1].unsqueeze(-1)
            logits = torch.where(logits < min_values, torch.tensor(float('-inf')).to(logits.device), logits)

        # 5. Top-p (核采样)过滤
        if top_p < 1.0:
            sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)

            # 移除累积概率超过top_p的token
            sorted_indices_to_remove = cumulative_probs > top_p
            # 确保至少保留一个token
            sorted_indices_to_remove[..., 0] = False

            # 创建散射掩码
            indices_to_remove = sorted_indices_to_remove.scatter(
                1, sorted_indices, sorted_indices_to_remove
            )
            logits = logits.masked_fill(indices_to_remove, float('-inf'))

        # 6. 从处理后的分布中采样
        probs = torch.softmax(logits, dim=-1)
        token_ids = torch.multinomial(probs, num_samples=1).squeeze(-1)

        return token_ids, kv_cache


# 创建模拟分词器
tokenizer = MockTokenizer(vocab_size=10000)

# 创建模型实例
config = PretrainedConfig(rms_norm_eps=1e-6)
rope_scaling = {"factor": 2.0}
model = DeepseekMLA_TextGeneration(
    config=config,
    hidden_size=4096,
    num_heads=32,
    qk_nope_head_dim=48,
    qk_rope_head_dim=16,
    v_head_dim=64,
    q_lora_rank=64,
    kv_lora_rank=64,
    rope_scaling=rope_scaling,
    vocab_size=tokenizer.vocab_size
).to(torch.float32)

# 模拟输入
prefill_seq_len = 10  # 使用更短的序列用于演示
hidden_states = torch.randn(prefill_seq_len, 4096, dtype=torch.float32)
positions = torch.arange(prefill_seq_len, dtype=torch.int64)
slot_mapping = torch.arange(prefill_seq_len, dtype=torch.int64)

# 生成文本
with torch.no_grad():
    token_ids, kv_cache = model.generate_text(
        hidden_states, positions, slot_mapping,
        temperature=0.8, top_k=50, top_p=0.9
    )

# 解码为文本
generated_text = tokenizer.decode(token_ids.tolist())
print("生成的Token IDs:", token_ids.tolist())
print("生成的文本:", generated_text)

# 演示如何继续生成更多文本 (解码阶段)
print("\n=== 模拟解码阶段 ===")
# 使用最后一个token作为下一个输入的起点
last_token_id = token_ids[-1]
print(f"最后一个Token ID: {last_token_id} ('{tokenizer.decode([last_token_id.item()], skip_special_tokens=False)}')")

# 在真实场景中，我们会:
# 1. 使用last_token_id作为下一个输入
# 2. 更新positions和slot_mapping
# 3. 调用模型的解码方法(如果有)或再次调用生成方法
# 4. 重复直到生成结束标记或达到最大长度

# 模拟继续生成
next_hidden_state = torch.randn(1, 4096, dtype=torch.float32)  # 模拟下一个隐藏状态
next_positions = torch.tensor([prefill_seq_len], dtype=torch.int64)  # 下一个位置
next_slot_mapping = torch.tensor([prefill_seq_len], dtype=torch.int64)  # 下一个缓存位置

# 在实际实现中，我们需要一个decode方法来处理单步生成
# 这里简化处理，直接使用generate_text
with torch.no_grad():
    next_token_ids, updated_kv_cache = model.generate_text(
        next_hidden_state, next_positions, next_slot_mapping,
        temperature=0.8, top_k=50, top_p=0.9
    )

print(
    f"下一个Token ID: {next_token_ids.item()} ('{tokenizer.decode([next_token_ids.item()], skip_special_tokens=False)}')")